{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qFdPvlXBOdUN"
      },
      "source": [
        "# Keras の例による量子化認識トレーニング"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://www.tensorflow.org/model_optimization/guide/quantization/training_example\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\"> TensorFlow.orgで表示</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/model_optimization/guide/quantization/training_example.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\"> Google Colab で実行</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ja/model_optimization/guide/quantization/training_example.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">GitHub でソースを表示</a></td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ja/model_optimization/guide/quantization/training_example.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\">ノートブックをダウンロード</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bjmi3qZeu_xk"
      },
      "source": [
        "## 概要\n",
        "\n",
        "*量子化認識トレーニング*のエンドツーエンドの例へようこそ。\n",
        "\n",
        "### その他のページ\n",
        "\n",
        "量子化認識トレーニングの紹介、および認識トレーニングを使用すべきかどうかの判定（サポート情報も含む）については、[概要ページ](https://www.tensorflow.org/model_optimization/guide/quantization/training.md)をご覧ください。\n",
        "\n",
        "ユースケースに合った API を素早く特定するには（8 ビットのモデルの完全量子化を超えるユースケース）、[総合ガイド](https://www.tensorflow.org/model_optimization/guide/quantization/training_comprehensive_guide.md)をご覧ください。\n",
        "\n",
        "### 要約\n",
        "\n",
        "このチュートリアルでは、次について説明しています。\n",
        "\n",
        "1. MNIST の `tf.keras` モデルを最初からトレーニングする。\n",
        "2. 量子化認識トレーニング API を適用してモデルをファインチューニングし、精度を確認して量子化認識モデルをエクスポートする。\n",
        "3. このモデルを使用して、TFLite バックエンドのために実際に量子化されたモデルを作成する。\n",
        "4. TFLite および  1/4 のモデルの精度の永続性を確認する。モバイルでのレイテンシーのメリットを確認するには、[TFLite アプリリポジトリ内の](https://www.tensorflow.org/lite/models) TFLite の例を試してみてください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yEAZYXvZU_XG"
      },
      "source": [
        "## セットアップ"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zN4yVFK5-0Bf"
      },
      "outputs": [],
      "source": [
        "! pip uninstall -y tensorflow\n",
        "! pip install -q tf-nightly\n",
        "! pip install -q tensorflow-model-optimization\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yJwIonXEVJo6"
      },
      "outputs": [],
      "source": [
        "import tempfile\n",
        "import os\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from tensorflow import keras"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "psViY5PRDurp"
      },
      "source": [
        "## 量子化認識トレーニングを使用せずに、MNIST のモデルをトレーニングする"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pbY-KGMPvbW9"
      },
      "outputs": [],
      "source": [
        "# Load MNIST dataset\n",
        "mnist = keras.datasets.mnist\n",
        "(train_images, train_labels), (test_images, test_labels) = mnist.load_data()\n",
        "\n",
        "# Normalize the input image so that each pixel value is between 0 to 1.\n",
        "train_images = train_images / 255.0\n",
        "test_images = test_images / 255.0\n",
        "\n",
        "# Define the model architecture.\n",
        "model = keras.Sequential([\n",
        "  keras.layers.InputLayer(input_shape=(28, 28)),\n",
        "  keras.layers.Reshape(target_shape=(28, 28, 1)),\n",
        "  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation='relu'),\n",
        "  keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
        "  keras.layers.Flatten(),\n",
        "  keras.layers.Dense(10)\n",
        "])\n",
        "\n",
        "# Train the digit classification model\n",
        "model.compile(optimizer='adam',\n",
        "              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "              metrics=['accuracy'])\n",
        "\n",
        "model.fit(\n",
        "  train_images,\n",
        "  train_labels,\n",
        "  epochs=1,\n",
        "  validation_split=0.1,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K8747K9OE72P"
      },
      "source": [
        "## 量子化認識トレーニングを使用して、事前トレーニング済みモデルをクローンおよびファインチューニングする\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F19k7ExXF_h2"
      },
      "source": [
        "### モデルを定義する"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JsZROpNYMWQ0"
      },
      "source": [
        "量子化認識トレーニングをモデル全体に適用し、これをモデルの要約で確認します。すべてのレイヤーにプレフィックス \"quant\" が付いているはずです。\n",
        "\n",
        "結果のモデルは量子化認識モデルですが、量子化はされていないので注意してください（例えば、重みは int8 ではなく float32 です）。次のセクションでは、量子化認識モデルから量子化モデルを作成する方法を示します。\n",
        "\n",
        "[総合ガイド](https://www.tensorflow.org/model_optimization/guide/quantization/training_comprehensive_guide.md)では、モデルの精度を改善するために、一部のレイヤーを量子化する方法をご覧いただけます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oq6blGjgFDCW"
      },
      "outputs": [],
      "source": [
        "import tensorflow_model_optimization as tfmot\n",
        "\n",
        "quantize_model = tfmot.quantization.keras.quantize_model\n",
        "\n",
        "# q_aware stands for for quantization aware.\n",
        "q_aware_model = quantize_model(model)\n",
        "\n",
        "# `quantize_model` requires a recompile.\n",
        "q_aware_model.compile(optimizer='adam',\n",
        "              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "              metrics=['accuracy'])\n",
        "\n",
        "q_aware_model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uDr2ijwpGCI-"
      },
      "source": [
        "### モデルをベースラインに対してトレーニングおよび評価する"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XUBEn94hXYB1"
      },
      "source": [
        "モデルを 1 エポックだけトレーニングした後のファインチューニングを実証するために、トレーニングデータのサブセットに対して、量子化認識トレーニングを使用してファインチューニングを行います。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_PHDGJryE31X"
      },
      "outputs": [],
      "source": [
        "train_images_subset = train_images[0:1000] # out of 60000\n",
        "train_labels_subset = train_labels[0:1000]\n",
        "\n",
        "q_aware_model.fit(train_images_subset, train_labels_subset,\n",
        "                  batch_size=500, epochs=1, validation_split=0.1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-byC2lYlMkfN"
      },
      "source": [
        "この例では、ベースラインと比較し、量子化認識トレーニング後のテスト精度の損失は、最小限あるいはゼロです。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6bMFTKSSHyyZ"
      },
      "outputs": [],
      "source": [
        "_, baseline_model_accuracy = model.evaluate(\n",
        "    test_images, test_labels, verbose=0)\n",
        "\n",
        "_, q_aware_model_accuracy = q_aware_model.evaluate(\n",
        "   test_images, test_labels, verbose=0)\n",
        "\n",
        "print('Baseline test accuracy:', baseline_model_accuracy)\n",
        "print('Quant test accuracy:', q_aware_model_accuracy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2IepmUPSITn6"
      },
      "source": [
        "## TFLite バックエンドの量子化モデルを作成する"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1FgNP4rbOLH8"
      },
      "source": [
        "この後に、重み int8 と活性化関数 uint8 を持つ、実際に量子化されたモデルが出来上がります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w7fztWsAOHTz"
      },
      "outputs": [],
      "source": [
        "converter = tf.lite.TFLiteConverter.from_keras_model(q_aware_model)\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "\n",
        "quantized_tflite_model = converter.convert()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BEYsyYVqNgeY"
      },
      "source": [
        "## TF から TFLite への精度の永続性を確認する"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "saadXD4JQsBK"
      },
      "source": [
        "テストデータセットで TFLite モデルを評価するヘルパー関数を定義します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b8yBouuGNqls"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "def evaluate_model(interpreter):\n",
        "  input_index = interpreter.get_input_details()[0][\"index\"]\n",
        "  output_index = interpreter.get_output_details()[0][\"index\"]\n",
        "\n",
        "  # Run predictions on every image in the \"test\" dataset.\n",
        "  prediction_digits = []\n",
        "  for i, test_image in enumerate(test_images):\n",
        "    if i % 1000 == 0:\n",
        "      print('Evaluated on {n} results so far.'.format(n=i))\n",
        "    # Pre-processing: add batch dimension and convert to float32 to match with\n",
        "    # the model's input data format.\n",
        "    test_image = np.expand_dims(test_image, axis=0).astype(np.float32)\n",
        "    interpreter.set_tensor(input_index, test_image)\n",
        "\n",
        "    # Run inference.\n",
        "    interpreter.invoke()\n",
        "\n",
        "    # Post-processing: remove batch dimension and find the digit with highest\n",
        "    # probability.\n",
        "    output = interpreter.tensor(output_index)\n",
        "    digit = np.argmax(output()[0])\n",
        "    prediction_digits.append(digit)\n",
        "\n",
        "  print('\\n')\n",
        "  # Compare prediction results with ground truth labels to calculate accuracy.\n",
        "  prediction_digits = np.array(prediction_digits)\n",
        "  accuracy = (prediction_digits == test_labels).mean()\n",
        "  return accuracy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TuEFS4CIQvUw"
      },
      "source": [
        "量子化されたモデルを評価し、TensorFlow の精度が TFLite バックエンドに持続されていることを確認します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VqQTyqz4NsWd"
      },
      "outputs": [],
      "source": [
        "interpreter = tf.lite.Interpreter(model_content=quantized_tflite_model)\n",
        "interpreter.allocate_tensors()\n",
        "\n",
        "test_accuracy = evaluate_model(interpreter)\n",
        "\n",
        "print('Quant TFLite test_accuracy:', test_accuracy)\n",
        "print('Quant TF test accuracy:', q_aware_model_accuracy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z8D7WnFF5DZR"
      },
      "source": [
        "## 量子化でモデルが 1/4 になることを確認する"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I1c2IecBRCdQ"
      },
      "source": [
        "浮動小数点数の TFLite モデルを作成して、量子化された TFLite モデルが 1/4 になっていることを確認します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jy_Lgfh8VkyX"
      },
      "outputs": [],
      "source": [
        "# Create float TFLite model.\n",
        "float_converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
        "float_tflite_model = float_converter.convert()\n",
        "\n",
        "# Measure sizes of models.\n",
        "_, float_file = tempfile.mkstemp('.tflite')\n",
        "_, quant_file = tempfile.mkstemp('.tflite')\n",
        "\n",
        "with open(quant_file, 'wb') as f:\n",
        "  f.write(quantized_tflite_model)\n",
        "\n",
        "with open(float_file, 'wb') as f:\n",
        "  f.write(float_tflite_model)\n",
        "\n",
        "print(\"Float model in Mb:\", os.path.getsize(float_file) / float(2**20))\n",
        "print(\"Quantized model in Mb:\", os.path.getsize(quant_file) / float(2**20))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0O5xuci-SonI"
      },
      "source": [
        "## 結論"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O2I7xmyMW5QY"
      },
      "source": [
        "このチュートリアルでは、TensorFlow Model Optimization Toolkit API を使用して量子化認識モデルを作成し、TFLite バックエンドの量子化モデルを作成する方法を紹介しました。\n",
        "\n",
        "MNIST のモデルでは、精度の違いを最小限に抑えながらモデルサイズを 1/4 に圧縮できることを示しましたた。モバイルでのレイテンシーのメリットを確認するには、[TFLite アプリリポジトリ内の](https://www.tensorflow.org/lite/models) TFLite の例を試してみてください。\n",
        "\n",
        "この新しい機能をぜひお試しください。リソースが制限される環境でのデプロイにおいて、特に重要となります。\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "training_example.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
